﻿#ifndef __DATA_H__
#define __DATA_H__

#include <cstdlib>
#include <ctime>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "Image.h"
#include "Options.h"

class InputOutput;
class Net;

class InputItem
{
public:
    InputItem(const std::string &image_path, const std::string &ofs_path) : m_image_path(image_path), m_ofs_path(ofs_path), m_count(0), m_input_index(0), m_image_index(0), m_ofs(nullptr), m_image(nullptr) {}

    inline void SetCount(int count) { m_count = count; }
    inline int GetCount() { return m_count; }
    inline void SetInputIndex(int input_index) { m_input_index = input_index; }
    inline int GetInputIndex() { return m_input_index; }
    inline void SetImageIndex(int image_index) { m_image_index = image_index; }
    inline int GetImageIndex() { return m_image_index; }
    inline void SetImage(std::shared_ptr<Image> image) { m_image = image; }
    inline std::shared_ptr<Image> GetImage() { return m_image; }
    inline void SetOFStream(std::shared_ptr<std::ofstream> ofs) { m_ofs = ofs; }
    inline std::ofstream *GetOFStream() { return (m_ofs != nullptr && m_ofs->good()) ? m_ofs.get() : nullptr; }
    inline std::vector<float> *GetTruths() { return &m_truths; }
    inline const std::string &GetOFSPath() { return m_ofs_path; }
    inline const std::string &GetImagePath() { return m_image_path; }

    std::shared_ptr<InputItem> Clone()
    {
        auto clone           = std::make_shared<InputItem>(m_image_path, m_ofs_path);
        clone->m_count       = m_count;
        clone->m_input_index = m_input_index;
        clone->m_image_index = m_image_index;
        clone->m_image       = m_image->Clone();
        clone->m_truths = m_truths;
        clone->m_ofs    = m_ofs;
        return clone;
    }

protected:
    std::string m_image_path;
    std::string m_ofs_path;
    int m_count;
    int m_input_index;
    int m_image_index;
    std::shared_ptr<Image> m_image;
    std::vector<float> m_truths;
    std::shared_ptr<std::ofstream> m_ofs;
};

class Input
{
public:
    Input(int width, int height, int channel, int batch, int truth_size)
    {
        m_width   = width;
        m_height  = height;
        m_channel = channel;
        m_batch   = batch;
        m_data.assign(m_width * m_height * m_channel * m_batch, 0);
        m_truth.assign(truth_size * m_batch, 0);
    }

public:
    inline float *GetTruth() { return m_truth.data(); }
    inline float *GetData() { return m_data.data(); }
    inline int GetWidth() { return m_width; }
    inline int GetHeight() { return m_height; }
    inline int GetChannel() { return m_channel; }
    inline int GetBatch() { return m_batch; }

public:
    std::vector<float> m_data;
    std::vector<float> m_truth;
    int m_width;
    int m_height;
    int m_channel;
    int m_batch;
};


class Data
{
public:
    Data(Net *net, const std::string &data_path = "");
    bool Load(const std::string &data_path);
    std::shared_ptr<Input> GetNextBatchTrainInput() { return GetNextBatchInput(&m_trains); }
    std::shared_ptr<Input> GetNextBatchTestInput() { return GetNextBatchInput(&m_tests); }
    std::vector<std::string> *GetTests() { return &m_tests; }
    InputItem *GetInputItem(int input_index);
    std::ofstream *GetInputOFS(int input_index);
    std::ofstream *GetTestOFS();
    std::ofstream *GetTrainOFS();

    void SaveNetOptionsToDefaultValues(Options *options);
    void LoadDefaultOptions();
    void LoadTrainOptions();

public:
    static void ImageBoxToNetBox(float *img_box, float *net_box, int img_width, int img_height, int net_width, int net_height);
    static void NetBoxToImageBox(float *net_box, float *img_box, int img_width, int img_height, int net_width, int net_height);
    static float GetNetImageScale(int img_width, int img_height, int net_width, int net_height, int *offset_x = nullptr, int *offset_y = nullptr);
    static std::shared_ptr<Image> LoadNetImage(const std::string &in_img_path, int in_net_width, int in_net_height, int *out_img_width = nullptr, int *out_img_height = nullptr);
    static std::shared_ptr<Image> LoadNetImage(const unsigned char* in_datas, const int in_size, int in_net_width, int in_net_height, int *out_img_width = nullptr, int *out_img_height = nullptr);
public:
    inline Net *GetNet() { return m_net; }
    inline bool IsLoaded() { return !m_path.empty(); }
    inline void SetFilpHorizontal(bool filp_horizontal) { m_filp_horizontal = filp_horizontal; }
    inline bool IsFilpHorizontal() { return m_filp_horizontal; }
    inline void SetFilpVertical(bool filp_vertical) { m_filp_vertical = filp_vertical; }
    inline bool IsFilpVertical() { return m_filp_vertical; }
    inline bool IsAutoStopTrain() { return m_auto_stop_train; }
    inline int GetMaxEpcohs() { return m_max_epochs; }

    inline std::string &GetNetPath() { return m_net_path; }
    inline std::string &GetModelsPath() { return m_models_path; }
    inline std::string &GetPretrainModelPath() { return m_pretrain_model_path; }
    inline std::vector<std::string> *GetNames() { return &m_names; }
    inline void SetInputIndex(int data_index) { m_input_index = data_index; }
    inline int GetInputIndex() { return m_input_index; }
    inline void SetRandom(bool random) { m_random = random; }
    inline bool IsRandom() { return m_random; }
    inline void SetDistort(bool distort) { m_distort = distort; }
    inline bool IsDistort() { return m_distort; }
    inline void SetJitter(float jitter) { m_jitter = jitter; }
    inline float GetJitter() { return m_jitter; }
    inline void SetDebug(bool debug) { m_debug = debug; }
    inline bool IsDebug() { return m_debug; }
    inline void SetHue(float hue) { m_hue = hue; }
    inline float GetHue() { return m_hue; }
    inline void SetSaturation(float saturation) { m_saturation = saturation; }
    inline float GetSaturation() { return m_saturation; }
    inline void SetExposure(float exposure) { m_exposure = exposure; }
    inline float GetExposure() { return m_exposure; }
    inline void SetBatch(int batch) { m_batch = batch; }
    inline int GetBatch() { return m_batch; }
    inline void SetSubdivisions(int subdivisions) { m_subdivisions = subdivisions; }
    inline int GetSubdivisions() { return m_subdivisions; }
    inline int GetPolicy() { return m_policy; }
    inline float GetMomentum() { return m_momentum; }
    inline float GetDecay() { return m_decay; }
    inline float GetLearningRate() { return m_learning_rate; }
    inline float* GetScales() { return m_scales.data(); }
    inline int* GetSteps() { return m_steps.data(); }
    inline int GetStepsSize() { return m_steps.size(); }
    inline float GetGamma() { return m_gamma; }
    inline float GetPower() { return m_power; }
    inline int GetEpochBatchSize() { return m_batch * m_subdivisions; }
protected:
    std::shared_ptr<Input> GetNextBatchInput(std::vector<std::string> *datas);
    bool LoadInputItemTruth(InputItem *input_item);
    std::shared_ptr<InputItem> AugmentInputItem(std::shared_ptr<InputItem> input_item);
    std::shared_ptr<InputItem> LoadInputItem(const std::string &in_path, const int truth_size, const int image_index, const int input_index);
    bool LoadLines(const std::string &in_path, std::vector<std::string> *out_lines);
    void LoadOptions(Options *options);

protected:
    Net *m_net;

    std::string m_names_path;
    std::string m_trains_path;
    std::string m_tests_path;
    std::vector<std::string> m_names;
    std::vector<std::string> m_trains;
    std::vector<std::string> m_tests;
    int m_input_index;
    std::unordered_map<int, std::shared_ptr<InputItem>> m_image_index_items;
    std::unordered_map<int, std::shared_ptr<InputItem>> m_input_index_items;
    std::string m_path;
    std::string m_net_path;
    std::string m_models_path;
    std::string m_debugs_path;
    std::string m_pretrain_model_path;

    int m_batch;                 // 训练批次
    int m_subdivisions;          // 训练子批次
    int m_max_epochs;            // 最大训练轮数
    float m_jitter;              // 抖动
    bool m_filp_horizontal;      // 是否水平翻转
    bool m_filp_vertical;        // 是否垂直翻转
    float m_saturation;          // 饱和度
    float m_exposure;            // 曝光
    float m_hue;                 // 色度
    bool m_distort;              // 是否变形图像
    bool m_random;               // 是否随机
    bool m_auto_stop_train;      // 是否自动停止训练
    bool m_debug;                // 是否调试
    int m_policy;                // 学习率策略
    float m_power;               // 学习率指数
    float m_gamma;               // 学习率指数
    std::vector<int> m_steps;    // 学习率调整点
    std::vector<float> m_scales; // 学习率调整点
    float m_learning_rate;       // 学习率
    float m_momentum;            // 增强
    float m_decay;               // 衰减

    std::shared_ptr<std::ofstream> m_test_ofs;
    std::shared_ptr<std::ofstream> m_train_ofs;

    std::shared_ptr<Options> m_default_options;
    std::shared_ptr<Options> m_train_options;
    std::shared_ptr<Options> m_valid_options;
    std::shared_ptr<Options> m_test_options;

protected:
    int m_default_batch;                 // 训练批次
    int m_default_subdivisions;          // 训练子批次
    int m_default_max_epochs;            // 最大训练轮数
    float m_default_jitter;              // 抖动
    bool m_default_filp_horizontal;      // 是否水平翻转
    bool m_default_filp_vertical;        // 是否垂直翻转
    float m_default_saturation;          // 饱和度
    float m_default_exposure;            // 曝光
    float m_default_hue;                 // 色度
    bool m_default_distort;              // 是否变形图像
    bool m_default_random;               // 是否随机
    bool m_default_auto_stop_train;      // 是否自动停止训练
    bool m_default_debug;                // 是否调试
    int m_default_policy;                // 学习率策略
    float m_default_power;               // 学习率指数
    float m_default_gamma;               // 学习率指数
    std::vector<int> m_default_steps;    // 学习率调整点
    std::vector<float> m_default_scales; // 学习率调整点
    float m_default_learning_rate;       // 学习率
    float m_default_momentum;            // 增强
    float m_default_decay;               // 衰减
protected:
    static const int s_default_image_channel = 3;
};

#endif